#!/usr/bin/env python3
"""
数据库模型定义
"""

from flask_sqlalchemy import SQLAlchemy
from datetime import datetime
import pytz
import bcrypt

db = SQLAlchemy()

# 获取东8区时间
def get_beijing_time():
    """获取东8区（北京时间）"""
    beijing_tz = pytz.timezone('Asia/Shanghai')
    return datetime.now(beijing_tz).replace(tzinfo=None)

class User(db.Model):
    """用户表"""
    __tablename__ = 'users'
    
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(120), unique=True, nullable=False)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联关系
    todos = db.relationship('Todo', backref='user', lazy=True, cascade='all, delete-orphan')
    agents = db.relationship('Agent', backref='user', lazy=True, cascade='all, delete-orphan')
    notes = db.relationship('Note', backref='user', lazy=True, cascade='all, delete-orphan')
    diary_entries = db.relationship('DiaryEntry', backref='user', lazy=True, cascade='all, delete-orphan')
    diary_summaries = db.relationship('DiarySummary', backref='user', lazy=True, cascade='all, delete-orphan')
    contacts = db.relationship('Contact', backref='user', lazy=True, cascade='all, delete-orphan')
    email_accounts = db.relationship('EmailAccount', backref='user', lazy=True, cascade='all, delete-orphan')
    email_messages = db.relationship('EmailMessage', backref='user', lazy=True, cascade='all, delete-orphan')
    prompt_configs = db.relationship('AIPromptConfig', backref='user', lazy=True, cascade='all, delete-orphan')
    
    def set_password(self, password):
        """设置密码（加密）"""
        self.password_hash = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
    
    def check_password(self, password):
        """验证密码"""
        return bcrypt.checkpw(password.encode('utf-8'), self.password_hash.encode('utf-8'))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'email': self.email,
            'username': self.username,
            'created_at': self.created_at.isoformat() if self.created_at else None,
        }


class Role(db.Model):
    """角色表"""
    __tablename__ = 'roles'
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(50), unique=True, nullable=False)
    description = db.Column(db.String(200))


class UserRole(db.Model):
    """用户-角色 关联表"""
    __tablename__ = 'user_roles'
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    role_id = db.Column(db.Integer, db.ForeignKey('roles.id'), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)


class Todo(db.Model):
    """待办事项表"""
    __tablename__ = 'todos'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    text = db.Column(db.Text, nullable=False)
    completed = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'text': self.text,
            'completed': self.completed,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class Agent(db.Model):
    """智能体表"""
    __tablename__ = 'agents'
    
    id = db.Column(db.String(100), primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    name = db.Column(db.String(100), nullable=False)
    url = db.Column(db.Text, nullable=False)
    description = db.Column(db.Text)
    icon = db.Column(db.String(50), default='MessageSquare')
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'name': self.name,
            'url': self.url,
            'description': self.description,
            'icon': self.icon,
            'createdAt': self.created_at.isoformat() if self.created_at else None,
            'updatedAt': self.updated_at.isoformat() if self.updated_at else None,
        }


class Note(db.Model):
    """Markdown笔记表"""
    __tablename__ = 'notes'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    title = db.Column(db.String(200), nullable=False)
    content = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'title': self.title,
            'content': self.content,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class DiaryEntry(db.Model):
    """日记条目表"""
    __tablename__ = 'diary_entries'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    content = db.Column(db.Text, nullable=False)
    mood = db.Column(db.String(20))  # 心情：开心、难过、平静等
    weather = db.Column(db.String(20))  # 天气：晴天、雨天、阴天等
    tags = db.Column(db.Text)  # 标签，用逗号分隔
    created_at = db.Column(db.DateTime, default=get_beijing_time)
    updated_at = db.Column(db.DateTime, default=get_beijing_time, onupdate=get_beijing_time)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'content': self.content,
            'mood': self.mood,
            'weather': self.weather,
            'tags': self.tags.split(',') if self.tags else [],
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class DiarySummary(db.Model):
    """AI日记总结表"""
    __tablename__ = 'diary_summaries'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    date = db.Column(db.Date, nullable=False)  # 总结的日期
    summary = db.Column(db.Text, nullable=False)  # AI生成的总结
    highlights = db.Column(db.Text)  # 亮点事件
    mood_analysis = db.Column(db.Text)  # 心情分析
    created_at = db.Column(db.DateTime, default=get_beijing_time)
    updated_at = db.Column(db.DateTime, default=get_beijing_time, onupdate=get_beijing_time)
    
    # 确保每个用户每天只有一个总结
    __table_args__ = (db.UniqueConstraint('user_id', 'date', name='unique_user_date_summary'),)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'date': self.date.isoformat() if self.date else None,
            'summary': self.summary,
            'highlights': self.highlights.split(',') if self.highlights else [],
            'mood_analysis': self.mood_analysis,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class Contact(db.Model):
    """通讯录联系人表"""
    __tablename__ = 'contacts'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    
    # 基本信息
    name = db.Column(db.String(100), nullable=False)  # 姓名
    phone = db.Column(db.String(20))  # 电话
    email = db.Column(db.String(120))  # 邮箱
    address = db.Column(db.Text)  # 地址
    
    # 扩展信息
    company = db.Column(db.String(100))  # 公司
    position = db.Column(db.String(100))  # 职位
    notes = db.Column(db.Text)  # 备注
    tags = db.Column(db.Text)  # 标签，用逗号分隔
    birthday = db.Column(db.Date)  # 生日
    
    # 管理字段
    group = db.Column(db.String(50), default='默认')  # 分组
    is_favorite = db.Column(db.Boolean, default=False)  # 收藏/星标
    avatar = db.Column(db.String(500))  # 头像URL
    social_accounts = db.Column(db.Text)  # 社交账号JSON字符串
    
    # 时间字段
    created_at = db.Column(db.DateTime, default=get_beijing_time)
    updated_at = db.Column(db.DateTime, default=get_beijing_time, onupdate=get_beijing_time)
    
    def to_dict(self):
        """转换为字典"""
        import json
        
        # 处理社交账号JSON
        social_accounts = {}
        if self.social_accounts:
            try:
                social_accounts = json.loads(self.social_accounts)
            except (json.JSONDecodeError, TypeError):
                social_accounts = {}
        
        return {
            'id': self.id,
            'name': self.name,
            'phone': self.phone,
            'email': self.email,
            'address': self.address,
            'company': self.company,
            'position': self.position,
            'notes': self.notes,
            'tags': self.tags.split(',') if self.tags else [],
            'birthday': self.birthday.isoformat() if self.birthday else None,
            'group': self.group,
            'is_favorite': self.is_favorite,
            'avatar': self.avatar,
            'social_accounts': social_accounts,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class EmailAccount(db.Model):
    """邮箱账号配置表"""
    __tablename__ = 'email_accounts'
    
    id = db.Column(db.String(100), primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    name = db.Column(db.String(100), nullable=False)  # 显示名称
    email = db.Column(db.String(120), nullable=False)  # 邮箱地址
    smtp_server = db.Column(db.String(100), nullable=False)  # SMTP服务器
    smtp_port = db.Column(db.Integer, nullable=False)  # SMTP端口
    imap_server = db.Column(db.String(100), nullable=False)  # IMAP服务器
    imap_port = db.Column(db.Integer, nullable=False)  # IMAP端口
    username = db.Column(db.String(120), nullable=False)  # 用户名
    password = db.Column(db.Text, nullable=False)  # 加密存储的密码
    use_ssl = db.Column(db.Boolean, default=True)  # 是否使用SSL
    is_default = db.Column(db.Boolean, default=False)  # 是否默认邮箱
    description = db.Column(db.Text)  # 描述
    last_fetch_at = db.Column(db.DateTime)  # 最后一次拉取邮件的时间
    is_fetching = db.Column(db.Boolean, default=False)  # 是否正在拉取邮件
    last_sync_uid = db.Column(db.Integer)  # 记录上次同步的最大UID
    is_first_fetch = db.Column(db.Boolean, default=True)  # 是否首次拉取
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联关系
    email_messages = db.relationship('EmailMessage', backref='email_account', lazy=True, cascade='all, delete-orphan')
    
    def to_dict(self):
        """转换为字典（不包含密码）"""
        return {
            'id': self.id,
            'name': self.name,
            'email': self.email,
            'smtp_server': self.smtp_server,
            'smtp_port': self.smtp_port,
            'imap_server': self.imap_server,
            'imap_port': self.imap_port,
            'username': self.username,
            'use_ssl': self.use_ssl,
            'is_default': self.is_default,
            'description': self.description,
            'last_fetch_at': self.last_fetch_at.isoformat() if self.last_fetch_at else None,
            'is_fetching': self.is_fetching,
            'last_sync_uid': self.last_sync_uid,
            'is_first_fetch': self.is_first_fetch,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class EmailMessage(db.Model):
    """邮件存储表"""
    __tablename__ = 'email_messages'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    email_account_id = db.Column(db.String(100), db.ForeignKey('email_accounts.id'), nullable=False)
    message_id = db.Column(db.String(200), nullable=False)  # 邮件唯一标识
    imap_uid = db.Column(db.Integer)  # IMAP服务器的唯一邮件ID
    subject = db.Column(db.String(500))  # 主题
    sender = db.Column(db.String(200))  # 发件人
    recipients = db.Column(db.Text)  # 收件人，JSON格式
    cc = db.Column(db.Text)  # 抄送，JSON格式
    body_text = db.Column(db.Text)  # 纯文本内容
    body_html = db.Column(db.Text)  # HTML内容
    attachments = db.Column(db.Text)  # 附件信息，JSON格式
    is_read = db.Column(db.Boolean, default=False)  # 是否已读
    is_starred = db.Column(db.Boolean, default=False)  # 是否星标
    is_full_content = db.Column(db.Boolean, default=False)  # 是否已加载完整正文
    content_size = db.Column(db.Integer)  # 邮件大小（字节）
    folder = db.Column(db.String(50), default='inbox')  # 文件夹：inbox/sent/draft/trash
    imap_folder = db.Column(db.String(200))  # IMAP原始文件夹名称
    received_at = db.Column(db.DateTime)  # 接收时间
    ai_summary = db.Column(db.Text)  # AI生成的邮件总结
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 索引
    __table_args__ = (
        db.Index('idx_message_id', 'message_id'),
        db.Index('idx_email_account_id', 'email_account_id'),
        db.Index('idx_folder', 'folder'),
        db.Index('idx_is_read', 'is_read'),
        db.Index('idx_imap_uid', 'imap_uid'),
        db.Index('idx_is_full_content', 'is_full_content'),
        # 新增: 复合索引优化查询
        db.Index('idx_user_folder_read', 'user_id', 'folder', 'is_read'),
        db.Index('idx_user_received_at', 'user_id', 'received_at'),
        db.Index('idx_user_account_folder', 'user_id', 'email_account_id', 'folder'),
        db.Index('idx_received_at_desc', 'received_at', postgresql_using='btree'),
        db.UniqueConstraint('user_id', 'message_id', name='unique_user_message'),
    )
    
    def to_dict(self):
        """转换为字典"""
        import json
        
        # 处理JSON字段
        recipients = []
        cc = []
        attachments = []
        
        if self.recipients:
            try:
                recipients = json.loads(self.recipients)
            except (json.JSONDecodeError, TypeError):
                recipients = []
        
        if self.cc:
            try:
                cc = json.loads(self.cc)
            except (json.JSONDecodeError, TypeError):
                cc = []
        
        if self.attachments:
            try:
                attachments = json.loads(self.attachments)
            except (json.JSONDecodeError, TypeError):
                attachments = []
        
        return {
            'id': self.id,
            'email_account_id': self.email_account_id,
            'message_id': self.message_id,
            'imap_uid': self.imap_uid,
            'subject': self.subject,
            'sender': self.sender,
            'recipients': recipients,
            'cc': cc,
            'body_text': self.body_text,
            'body_html': self.body_html,
            'attachments': attachments,
            'is_read': self.is_read,
            'is_starred': self.is_starred,
            'is_full_content': self.is_full_content,
            'content_size': self.content_size,
            'folder': self.folder,
            'imap_folder': self.imap_folder,
            'received_at': self.received_at.isoformat() if self.received_at else None,
            'ai_summary': self.ai_summary,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }


class AIPromptConfig(db.Model):
    """AI提示词配置表"""
    __tablename__ = 'ai_prompt_configs'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    module = db.Column(db.String(50), nullable=False)  # 'email', 'chat', 'notes' 等
    prompt_template = db.Column(db.Text, nullable=False)
    is_active = db.Column(db.Boolean, default=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 添加唯一约束：每个用户的每个模块只能有一个激活的配置
    __table_args__ = (
        db.UniqueConstraint('user_id', 'module', name='uq_user_module'),
    )
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'user_id': self.user_id,
            'module': self.module,
            'prompt_template': self.prompt_template,
            'is_active': self.is_active,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
        }

